package cn.cas.mango.service;

import cn.cas.mango.dto.param.AggParam;
import cn.cas.mango.dto.param.News;
import cn.cas.mango.util.Pair;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.FieldValue;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregate;
import co.elastic.clients.elasticsearch._types.aggregations.BucketMetricValueAggregate;
import co.elastic.clients.elasticsearch._types.aggregations.CompositeAggregationSource;
import co.elastic.clients.elasticsearch._types.aggregations.CompositeBucket;
import co.elastic.clients.elasticsearch._types.aggregations.StringTermsBucket;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

@Service
@Slf4j
public class Agg {
  @Value("${elasticsearch.newIndex}")
  private String index;
  private String field = "news_cluster_id";
  @Autowired
  private ElasticsearchClient client;
  public JSONObject agg(AggParam aggParam) {
    SearchRequest clusterGroup = SearchRequest.of(x -> x.index(index).trackTotalHits(k->k.enabled(true)).aggregations("cluster_group", y -> y.terms(m -> m.field(field).size(aggParam.getSize()))));
    JSONObject jsonObject = new JSONObject();


    try {
      SearchResponse<Map> search = client.search(clusterGroup, Map.class);
      Collection<Aggregate> values = search.aggregations().values();
      JSONArray jsonArray = new JSONArray();
      for (Aggregate aggregate: values) {
        List<StringTermsBucket> array = aggregate.sterms().buckets().array();
        for (StringTermsBucket stringTermsBucket: array) {
          JSONObject jsonObject1 = new JSONObject();
          jsonObject1.put((String) stringTermsBucket.key()._get(), stringTermsBucket.docCount());
          jsonArray.add(jsonObject1);
        }
      }
      jsonObject.put("totalCounts", jsonArray.size());
      jsonObject.put("data", jsonArray);
      return
          jsonObject;
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  public List<Hit<News>> getOne(AggParam aggParam) {
    SearchRequest of = SearchRequest.of(x -> x.index(index)
        .query(q -> q.term(t -> t.field(field)
            .value(aggParam.getNewsClusterId()))).size(aggParam.getSize() == null ? 100: aggParam.getSize()));
    try {
      return client.search(of, News.class).hits().hits();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  public List<Hit<News>> getAnother(AggParam aggParam) {
    SearchRequest of = SearchRequest.of(x -> x.index(index)
        .query(q -> q.term(t -> t.field(field)
            .value(aggParam.getNewsClusterId()))).size(aggParam.getSize() == null ? 100: aggParam.getSize()).source(
                s->s.filter(v->v.includes("id", "news_title_zh", "news_content_zh"))
        ));
    try {
      return client.search(of, News.class).hits().hits();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  public void agg(List<String> ret, String key) {
    SearchRequest searchRequest;
    if (key == null )
      searchRequest = SearchRequest.of(x -> x.index(index).aggregations("temp", a -> a.composite(c -> c.sources(
        new HashMap<String, CompositeAggregationSource>() {{
          put("cluster", CompositeAggregationSource.of(f -> f.terms(t -> t.field(field))));
        }}).size(10000)
        ))
    );
    else {
      searchRequest = SearchRequest.of(x -> x.index(index).aggregations("temp", a -> a.composite(c -> c.sources(
                  new HashMap<String, CompositeAggregationSource>() {{
                    put("cluster", CompositeAggregationSource.of(f -> f.terms(t -> t.field(field))));
                  }}).size(10000)
              .after("cluster", FieldValue.of(key))
          ))
      );
    }
//    System.out.println(searchRequest);
    try {
      SearchResponse<News> search = client.search(searchRequest, News.class);
      Map<String, Aggregate> aggregations = search.aggregations();
      Collection<Aggregate> values = aggregations.values();
      for (Aggregate aggregate: values) {
        List<CompositeBucket> array1 = aggregate.composite().buckets().array();

        for (CompositeBucket c: array1) {
          String s = c.key().get("cluster").stringValue();
          ret.add(s);
        }
        if (!aggregate.composite().afterKey().isEmpty()) {
          agg(ret, aggregate.composite().afterKey().get("cluster").stringValue());
        }
      }
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  public JSONArray get(AggParam aggParam) {
    JSONArray ret = new JSONArray();

    List<String> ask = new ArrayList<>();
    agg(ask, null);
    int count = 0;
    List<Pair<String, List<Hit<News>>>> list = new ArrayList<>();

    for (String string: ask) {
      aggParam.setNewsClusterId(string);
      List<Hit<News>> one = getAnother(aggParam);
      list.add(new Pair<>(string, one));

      count ++;
      if (count % 1000 == 0) {
        log.info("finished "+ count + "/"+ ask.size());
      }
    }
    list.sort((x,y)->{
      if (x.getValue().size() == y.getValue().size()) {
        return 0;
      }
      return x.getValue().size() > y.getValue().size() ? -1: 1;
    });
    for (Pair<String, List<Hit<News>>>l : list) {
      String string  = l.getKey();
      List<Hit<News>> one = l.getValue();
      JSONObject temp = new JSONObject();
      temp.put("clusterId", string);
      temp.put("members",
          one.stream().map(x->
          {JSONObject t = new JSONObject();
            News source = x.source();
            t.put("id", source.getId());
            t.put("title", source.getNewsTitleZh());
            t.put("content", source.getNewsContentZh());
            return t;
          }).collect(Collectors.toList()));
      ret.add(temp);
    }
    return ret;
  }


}
